// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/exec/aggregate.h>
#include <pollux/expression/function_signature.h>
#include <pollux/functions/lib/aggregates/simple_numeric_aggregate.h>

namespace kumo::pollux::functions::aggregate {
    template<typename T>
    class BitwiseAggregateBase : public SimpleNumericAggregate<T, T, T> {
        using BaseAggregate = SimpleNumericAggregate<T, T, T>;

    public:
        BitwiseAggregateBase(TypePtr resultType, T initialValue)
            : BaseAggregate(resultType), initialValue_(initialValue) {
        }

        int32_t accumulatorFixedWidthSize() const override {
            return sizeof(T);
        }

        void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
        override {
            BaseAggregate::doExtractValues(groups, numGroups, result, [&](char *group) {
                return *BaseAggregate::Aggregate::template value<T>(group);
            });
        }

        void addIntermediateResults(
            char **groups,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool mayPushdown) override {
            this->addRawInput(groups, rows, args, mayPushdown);
        }

        void addSingleGroupIntermediateResults(
            char *group,
            const SelectivityVector &rows,
            const std::vector<VectorPtr> &args,
            bool mayPushdown) override {
            this->addSingleGroupRawInput(group, rows, args, mayPushdown);
        }

    protected:
        void initializeNewGroupsInternal(
            char **groups,
            melon::Range<const vector_size_t *> indices) override {
            exec::Aggregate::setAllNulls(groups, indices);
            for (auto i: indices) {
                *exec::Aggregate::value<T>(groups[i]) = initialValue_;
            }
        }

        const T initialValue_;
    };

    template<template <typename U> class T>
    exec::AggregateRegistrationResult registerBitwise(
        const std::string &name,
        bool withCompanionFunctions,
        bool onlyPrestoSignatures,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures;
        std::vector<std::string> typeList{"tinyint", "smallint", "integer", "bigint"};
        if (onlyPrestoSignatures) {
            typeList = {"bigint"};
        }
        for (const auto &inputType: typeList) {
            signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                .returnType(inputType)
                .intermediateType(inputType)
                .argumentType(inputType)
                .build());
        }

        return exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step /*step*/,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig &config) -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_LE(argTypes.size(), 1, "{} takes only one argument", name);
                auto inputType = argTypes[0];
                switch (inputType->kind()) {
                    case TypeKind::TINYINT:
                        return std::make_unique<T<int8_t> >(resultType);
                    case TypeKind::SMALLINT:
                        return std::make_unique<T<int16_t> >(resultType);
                    case TypeKind::INTEGER:
                        return std::make_unique<T<int32_t> >(resultType);
                    case TypeKind::BIGINT:
                        return std::make_unique<T<int64_t> >(resultType);
                    default:
                        POLLUX_UNREACHABLE(
                            "Unknown input type for {} aggregation {}",
                            name,
                            inputType->kindName());
                }
            },
            {false /*orderSensitive*/, false /*companionFunction*/},
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::functions::aggregate
